import numpy as np

from .base import BaseAMM

from market import Market


class LMSR(BaseAMM):
    def __init__(self, b: float):
        self.b = b

    def _price(self, true_share: int, false_share: int) -> float:
        return 1 / (1 + np.exp((false_share - true_share) / self.b))

    def probability_after_trade(self, market: Market, is_true: bool, share: int) -> float:
        if is_true:
            return self._price(market.true_share + share, market.false_share)

        return self._price(market.true_share, market.false_share + share)

    def trade(self, market: Market, is_true: bool, share: float):
        if is_true:
            market.true_share += share
        else:
            market.false_share += share

    def price(self, market: Market) -> float:
        return self._price(market.true_share, market.false_share)

    def probability(self, market: Market) -> float:
        """
        LMSR中，选项价格就是选项概率
        """
        return self.price(market)
